# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
from matplotlib import cm
import pickle
import pandas as pd


def location_diagram(RA, DE, BV):
    '''
    画出 位置图，其中颜色根据色指数
    '''
    
    fig, ax = plt.subplots()
    cmap = cm.hot
    ax.scatter(RA, DE, c=BV, cmap=cmap)
    ax.set_xlabel(r'RA', fontsize=15)
    ax.set_ylabel(r'DE', fontsize=15)
    ax.set_title('location diagram')
    ax.grid(True)
    ax.set_facecolor((0, 0, 0))
    fig.colorbar(cm.ScalarMappable(cmap=cmap))
    plt.show()

def speed_diagram(pmRA, pmDE, BV):
    '''
    画出 速度图
    '''
    
    fig, ax = plt.subplots()
    cmap = cm.hot
    ax.scatter(pmRA, pmDE, c=BV, cmap=cmap)
    ax.set_xlabel(r'pmRA', fontsize=15)
    ax.set_ylabel(r'pmDE', fontsize=15)
    ax.set_title('speed diagram')
    ax.grid(True)
    ax.set_facecolor((0, 0, 0))
    fig.colorbar(cm.ScalarMappable(cmap=cmap))
    plt.show()

if __name__ == '__main__':
    bi_stream = pickle.load(open(r'../附件/bi_stream.pkl', 'rb'))
    print(bi_stream)
    
    RA = bi_stream['RA']
    DE = bi_stream['DE']
    pmRA = bi_stream['pmRA']
    pmDE = bi_stream['pmDE']
    BV = bi_stream['B-V']
    # 描述毕宿星流
    location_diagram(RA, DE, BV)
    speed_diagram(pmRA, pmDE, BV)
    
    bi_clu_label = pickle.load(open(r'../附件/bi_star_cluster_label.pkl', 'rb'))
    cluster_res = pickle.load(open(r'../附件/clustering_result.pkl', 'rb'))
    path = r'../附件/星表数据.txt'
    # 读取 txt 文件
    df = pd.read_csv(path, delimiter=' +')
    df['label'] = cluster_res
    # 找出 【20,22】 的点
    Plx_20_22 = df.loc[(20<=df['Plx']) & (df['Plx']<=22)].index
    df = df.iloc[Plx_20_22]
    # 提取毕星团
    bi_stars = df.loc[df['label']==bi_clu_label]
    RA = bi_stars['RA']
    DE = bi_stars['DE']
    pmRA = bi_stars['pmRA']
    pmDE = bi_stars['pmDE']
    BV = bi_stars['B-V']
    
    
    # 描述毕星团
    location_diagram(RA, DE, BV)
    speed_diagram(pmRA, pmDE, BV)
